import numpy as np


class MyDataSet:
    def __init__(self, all_datas, batch_size, shuffle=True):
        self.all_datas = all_datas
        self.batch_size = batch_size
        self.shuffle = shuffle

    def __iter__(self):  # 循环时候自动触发，返回一个具有__next__方法的对象。用来循环DataSet
        return MyDataLoader(self)

    def __len__(self):
        return len(self.all_datas)


class MyDataLoader:
    def __init__(self, dataset):
        self.dataSet = dataset
        self.indexs = [i for i in range(len(self.dataSet))]  # 用到了上面的__len__
        if self.dataSet.shuffle:
            np.random.shuffle(self.indexs)
        self.cursor = 0

    def __next__(self):  # 循环完一个位置到下一个位置。
        # 如果到头了，就不用到下一个了，抛出异常退出
        if self.cursor >= len(self.dataSet.all_datas):
            raise StopIteration

        # 不然就取出batch_size大小的数据
        index = self.indexs[self.cursor: self.cursor + self.dataSet.batch_size]

        # 取出index数组里对应的下标
        batch_data = self.dataSet.all_datas[index]
        self.cursor += self.dataSet.batch_size
        return batch_data


def test():
    all_datas = np.array([1, 2, 3, 4, 5, 6, 7])
    batch_size = 2
    shuffle = True
    epoch = 2

    dataset = MyDataSet(all_datas, batch_size, shuffle)

    for e in range(epoch):
        for batch_data in dataset:  # 把一个对象放在for上时, 会自动调用这个对象的__iter__,
            print(batch_data)


test()
